package xyz.scootaloo.code.adt;

import org.jetbrains.annotations.NotNull;

/**
 * @author flutterdash@qq.com
 * @since 2021/11/11 15:55
 */
public class AvlTree<E extends Comparable<E>>
        extends AbstractBinarySearchTree<E> implements KCollection<E> {

    private static final int ALLOW_IMBALANCE = 1;

    private int size;
    private AvlNode<E> root;

    public AvlTree() {
        clear();
    }

    @Override
    public int size() {
        return this.size;
    }

    @Override
    public void clear() {
        root = null;
        size = 0;
    }

    @Override
    public boolean add(@NotNull E element) {
        if (contains(element))
            return false;
        root = insertNode(root, element);
        size++;
        return true;
    }

    @Override
    public boolean remove(@NotNull E element) {
        if (!contains(element))
            return false;
        root = removeNode(root, element);
        size--;
        return true;
    }

    @Override
    public boolean contains(@NotNull E element) {
        AvlNode<E> current = root;
        int diff;
        while (current != null) {
            diff = current.value.compareTo(element);
            if (diff < 0) {
                current = current.right;
            } else if (diff > 0) {
                current = current.left;
            } else {
                return true;
            }
        }
        return false;
    }

    @NotNull
    @Override
    public KIterator<E> iterator() {
        return treeIterator(root);
    }

    @Override
    public boolean isEmpty() {
        return size != 0;
    }

    private AvlNode<E> removeNode(AvlNode<E> current, E element) {
        if (current == null)
            return null;

        int diff = current.value.compareTo(element);
        if (diff < 0) {
            current.left = removeNode(current.left, element);
        } else if (diff > 0) {
            current.right = removeNode(current.right, element);
        } else if (current.left != null && current.right != null) {
            current.value = findMin(current.right).value;
            current.right = removeNode(current.right, current.value);
        } else {
            current = (current.left != null) ? current.left : current.right;
        }

        return balance(current);
    }

    private AvlNode<E> insertNode(AvlNode<E> current, E element) {
        if (current == null)
            return createNode(element);

        int diff = current.value.compareTo(element);
        if (diff < 0) {
            current.right = insertNode(current.right, element);
        } else if (diff > 0) {
            current.left = insertNode(current.left, element);
        }

        current.updateDepth();
        return balance(current);
    }

    /**
     * 再平衡
     *
     * 1. 左子树深度大于右子树
     * 2. 右子树深度大于左子树
     *
     * @param current 执行再平衡操作的树节点
     * @return        执行再平衡操作后的实际根节点
     */
    private AvlNode<E> balance(AvlNode<E> current) {
        if (current == null)
            return null;

        int factor = treeBalanceFactor(current);
        if (factor > ALLOW_IMBALANCE) {
            if (treeBalanceFactor(current.left) > 0)
                current = rightRotate(current);
            else
                current = doubleWithLeftChild(current);
        } else if (factor < -ALLOW_IMBALANCE) {
            if (treeBalanceFactor(current.right) <= 0)
                current = leftRotate(current);
            else
                current = doubleWithRightChild(current);
        }

        current.updateDepth();
        return current;
    }

    private AvlNode<E> doubleWithLeftChild(AvlNode<E> current) {
        current.left = rightRotate(current.left);
        return leftRotate(current);
    }

    private AvlNode<E> doubleWithRightChild(AvlNode<E> current) {
        current.right = leftRotate(current);
        return rightRotate(current);
    }

    /**
     * 右旋一个树节点
     *
     *   右旋前:         右旋后:
     *       A              B
     *     /   \          /   \
     *    B    (AR)     C      A
     *  /  \                 /   \
     * C  (BR)             (BR)  (AR)
     *
     * AR: 执行旋转操作前根节点A的右孩子
     * BR: 执行旋转操作前节点B的右孩子
     *
     * 1. 右旋转前提条件, 以A为根节点的树, 其左子树高度大于右子数高度;
     * 2. 根据1得出, A的左孩子节点必定不为空, 因为要使左子数高度高于右子数, 右子树高度最少取0, 左子数高度最少取1;
     * 3. 选取节点B做为根节点;
     * 4. B的右孩子(BR)的值大于B而小于A, 所以(BR)可以做为A的左孩子(A的左孩子原来是B, 现在替换为(BR));
     * 5. B的右孩子(BR)现在替换为A;
     * 6. 更新节点(BR)和A的父节点信息;
     * 7. 更新以B为根节点的所有孩子节点的高度;
     * 8. 返回根节点B;
     *
     * @param current 要执行右旋转操作的树节点
     * @return        执行旋转操作后的根节点
     */
    private AvlNode<E> rightRotate(AvlNode<E> current) {
        if (current == null || current.left == null) // 2: 左子数为空, 无法执行右旋操作
            return current;

        AvlNode<E> vRoot = current.left; // 3: vRoot = B

        current.left = vRoot.right; // 4: A.lChild = B.rChild
        vRoot.right = current; // 5: B.rChild = A

        return vRoot; // 8
    }

    /**
     * 左旋一个树节点
     *
     * 左旋前:            左旋后:
     *     A                 B
     *   /   \             /   \
     * (AL)   B           A     C
     *      /  \        /   \
     *   (BL)   C     (AL) (BL)
     *
     * 1,2,3,4,5,6,7,8
     * 左旋操作实现细节参考右旋
     *
     * @param current 要执行左旋操作的节点
     * @return        旋转后最终得到的根节点
     */
    private AvlNode<E> leftRotate(AvlNode<E> current) {
        if (current == null || current.right == null) // 2
            return current;

        AvlNode<E> vRoot = current.right; // 3

        current.right = vRoot.left; // 4
        vRoot.left = current; // 5

        return vRoot; // 8
    }

    private int treeBalanceFactor(AvlNode<E> node) {
        if (node == null) return 0;
        return (node.left != null ? node.left.depth : -1) -
                (node.right != null ? node.right.depth : -1);
    }

    private AvlNode<E> createNode(E element) {
        return new AvlNode<>(element);
    }

    private static class AvlNode<E extends Comparable<E>> extends BinaryNode<E> {

        public int depth = 0;
        public AvlNode<E> left;
        public AvlNode<E> right;

        public AvlNode(E value) {
            super(value);
        }

        public int updateDepth() {
            this.depth = Math.max(
                    this.left != null ? this.left.updateDepth() : -1,
                    this.right != null ? this.right.updateDepth() : -1
            ) + 1;
            return this.depth;
        }

        @Override
        public String toString() {
            return String.format("Avl(V=%s, L=%s, R=%s, D=%d)", value,
                    (left == null ? "null" : left.value),
                    (right == null ? "null" : right.value), depth);
        }
    }

}
